import genesis as gs
import numpy as np
import torch

gs.init(backend=gs.gpu)

viewer_options = gs.options.ViewerOptions(
    camera_pos=(0, -3.5, 2.5),
    camera_lookat=(0.0, 0.0, 0.5),
    camera_fov=40,
    max_FPS=30,
)

scene = gs.Scene(
    viewer_options=viewer_options,
    sim_options=gs.options.SimOptions(
        dt=0.01,
        substeps=1,
        gravity=(0, 0, -10),
    ),
    show_viewer=True,
    renderer=gs.renderers.Rasterizer(),
)

ground = scene.add_entity(
    gs.morphs.Plane()
)

franka = scene.add_entity(
    gs.morphs.MJCF(
        file="xml/franka_emika_panda/panda.xml",
        pos = (0, 0, 0),
        euler = np.random.rand(3) * 2 * np.pi,
        scale = 1.0,          
    ),
)

sphere_pos = (np.random.rand(), np.random.rand(), 0)

target_sphere = scene.add_entity(
    gs.morphs.Sphere(
        radius=0.05,
        pos=sphere_pos
    ),
)

batches = 128
scene.build(n_envs=batches, env_spacing=(2.0, 2.0))

joint_names = [
    "joint1",
    "joint2",
    "joint3",
    "joint4",
    "joint5",
    "joint6",
    "joint7",
    "finger_joint1",
    "finger_joint2",
]

dofs_idx = [franka.get_joint(name).dof_idx_local for name in joint_names]



############ 可选：设置控制增益 ############
# 设置位置增益
franka.set_dofs_kp(
    kp             = np.array([4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100]),
    dofs_idx_local = dofs_idx,
)
# 设置速度增益
franka.set_dofs_kv(
    kv             = np.array([450, 450, 350, 350, 200, 200, 200, 10, 10]),
    dofs_idx_local = dofs_idx,
)
# 设置安全的力范围
franka.set_dofs_force_range(
    lower          = np.array([-87, -87, -87, -87, -12, -12, -12, -100, -100]),
    upper          = np.array([ 87,  87,  87,  87,  12,  12,  12,  100,  100]),
    dofs_idx_local = dofs_idx,
)


franka.control_dofs_position(
    torch.zeros(batches, 9, device=gs.device),
    dofs_idx,
)
# PD控制

total_steps = 1000
while True:
    franka.control_dofs_position(
        torch.randn(batches, 9, device=gs.device),
        dofs_idx,
    )

    # print(f"Step control force: ", franka.get_dofs_force(dofs_idx))
    # print(f"Step inner force: ", franka.get_dofs_force(dofs_idx))
    scene.step()